package com.witmore.neutron.ai.cognitive.service;

import com.alibaba.fastjson2.JSONObject;
import com.baomidou.mybatisplus.core.conditions.query.LambdaQueryWrapper;
import com.witmore.neutron.ai.cognitive.mcp.config.WitmoreSyncMcpToolCallbackProvider;
import com.witmore.neutron.ai.db.entity.AgentDO;
import com.witmore.neutron.ai.db.entity.AgentPluginDO;
import com.witmore.neutron.ai.db.entity.AgentSessionRecordDO;
import com.witmore.neutron.ai.db.entity.PluginDO;
import com.witmore.neutron.ai.db.mapper.AgentMapper;
import com.witmore.neutron.ai.db.mapper.AgentPluginMapper;
import com.witmore.neutron.ai.db.mapper.AgentSessionRecordMapper;
import com.witmore.neutron.ai.db.mapper.PluginMapper;
import lombok.extern.slf4j.Slf4j;
import org.springframework.ai.tool.ToolCallback;
import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.stereotype.Service;
import org.springframework.util.CollectionUtils;

import java.util.Arrays;
import java.util.List;
import java.util.Map;
import java.util.stream.Collectors;

/**
 * Description: PluginServiceImpl
 * Author: 豆子高
 * Date: 2025/7/4 10:47
 */
@Service
@Slf4j
public class AgentPluginServiceImpl implements AgentPluginService {

    @Autowired
    private AgentSessionRecordMapper agentSessionRecordMapper;

    @Autowired
    private AgentMapper agentMapper;

    @Autowired
    private AgentPluginMapper agentPluginMapper;

    @Autowired
    private PluginMapper pluginMapper;

    @Autowired
    private WitmoreSyncMcpToolCallbackProvider toolCallbackProvider;

    @Override
    public ToolCallback[] getToolCallbacks(String sessionId) {
        AgentSessionRecordDO agentSessionRecordDO = agentSessionRecordMapper.selectOne(
                new LambdaQueryWrapper<AgentSessionRecordDO>()
                        .eq(AgentSessionRecordDO::getSessionId, sessionId)
        );
        if (agentSessionRecordDO == null) {
            return null;
        }

        AgentDO agentDO = agentMapper.selectById(agentSessionRecordDO.getAgentId());
        if (agentDO == null) {
            return null;
        }

        // 获取智能体绑定的插件组列表
        List<AgentPluginDO> agentPluginGroups = agentPluginMapper.selectList(
                new LambdaQueryWrapper<AgentPluginDO>()
                        .eq(AgentPluginDO::getAgentId, agentDO.getId())
                        .eq(AgentPluginDO::getEnabled, 1)
        );

        if (CollectionUtils.isEmpty(agentPluginGroups)) {
            return null;
        }

        // 获取插件组下的所有插件
        List<String> groupCodes = agentPluginGroups.stream()
                .map(AgentPluginDO::getGroupCode)
                .collect(Collectors.toList());

        List<PluginDO> plugins = pluginMapper.selectList(
                new LambdaQueryWrapper<PluginDO>()
                        .in(PluginDO::getGroupCode, groupCodes)
                        .eq(PluginDO::getEnabled, 1)
                        .orderByAsc(PluginDO::getSortOrder)
        );

        if (CollectionUtils.isEmpty(plugins)) {
            return null;
        }

        // 将插件转换成Map，以pluginCode为key
        Map<String, PluginDO> pluginMap = plugins.stream()
                .collect(Collectors.toMap(PluginDO::getPluginCode, plugin -> plugin));

        if (CollectionUtils.isEmpty(pluginMap)) {
            return null;
        }

        // 获取MCP Server远程插件库 (使用新的降级机制)
        ToolCallback[] toolCallbacks;
        try {
            toolCallbacks = toolCallbackProvider.getToolCallbacks();
            if (toolCallbacks.length == 0) {
                log.warn("无法获取MCP服务工具回调，会话：{}，但不影响主流程，继续使用大模型基础功能", sessionId);
                return null; // 返回null，让大模型在没有工具的情况下正常工作
            }
        } catch (Exception e) {
            log.error("获取MCP工具回调时发生异常，会话：{}，错误信息：{}，不影响主流程，继续使用大模型基础功能", sessionId, e.getMessage());
            return null; // 异常时也返回null，确保主流程不受影响
        }

        // 过滤出智能体可使用的工具
        try {
            ToolCallback[] array = Arrays.stream(toolCallbacks)
                    .filter(toolCallback -> {
                        try {
                            boolean b = pluginMap.containsKey(toolCallback.getToolDefinition().name());
                            return b;
                        } catch (Exception e) {
                            log.warn("检查工具定义时发生异常：{}，跳过该工具", e.getMessage());
                            return false; // 跳过有问题的工具
                        }
                    })
                    .toArray(ToolCallback[]::new);

            if (array.length == 0) {
                log.warn("当前会话:{} 没有可用的工具，但不影响主流程，继续使用大模型基础功能", sessionId);
                return null;
            }

            log.info("当前会话:{} 开启的插件有：{}", sessionId, JSONObject.toJSONString(array));
            return array;
        } catch (Exception e) {
            log.error("过滤工具时发生异常，会话：{}，错误信息：{}，不影响主流程，继续使用大模型基础功能", sessionId, e.getMessage());
            return null;
        }
    }


}
